#include <math.h>
#include <complex.h>
#include <stdlib.h>
#include <stdio.h>
#include <drfftw_mpi.h>
#include <mpi.h>
#include <gsl/gsl_rng.h>

#include "allvars.h"
#include "proto.h"

#define sqr(x) ((x)*(x))

#define ASSERT_ALLOC(cond) {                                                                                  \
   if(cond)                                                                                                   \
    {                                                                                                         \
      if(ThisTask == 0)                                                                                       \
	printf("\nallocated %g Mbyte on Task %d\n", bytes / (1024.0 * 1024.0), ThisTask);                     \
    }                                                                                                         \
  else                                                                                                        \
    {                                                                                                         \
      printf("failed to allocate %g Mbyte on Task %d\n", bytes / (1024.0 * 1024.0), ThisTask);                \
      printf("bailing out.\n");                                                                               \
      FatalError(1);                                                                                          \
    }                                                                                                         \
}
//float dot(float x[3], float y[3]);
float dot(float x[3], float y[3]) {return x[0]*y[0]+x[1]*y[1]+x[2]*y[2];}

void statistics(fftw_complex *c1, fftw_complex *c2, float a_init, float xsize, long n1, float R);

fftw_complex ***fftw_ctensor3(int n1, int n2, int n3)
{
  int i, j;
  fftw_complex ***m;
  m=(fftw_complex ***)malloc(n1*sizeof(fftw_complex **));
  *m=(fftw_complex **)malloc(n1*n2*sizeof(fftw_complex *));
  if (!(**m=(fftw_complex *)calloc(n1*n2*(n3/2 + 1),sizeof(fftw_complex)))) {
    fprintf(stderr, "Error: unable to allocate %d*%d*%d ctensor\n",n1,n2,n3);
    perror("calloc");
    return NULL;
  }
  for(i=0;i<n1;i++) {
    if (i) {m[i]=m[i-1]+n2; *(m[i]) = *(m[i-1]) + n2*(n3/2 + 1);}
    for (j=1; j<n2; j++) m[i][j] = m[i][j-1]+(n3/2 + 1);
  }
  return m;
}

int main(int argc, char **argv)
{
  MPI_Init(&argc, &argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &ThisTask);
  MPI_Comm_size(MPI_COMM_WORLD, &NTask);

  if(argc < 2)
    {
      if(ThisTask == 0)
	{
	  fprintf(stdout, "\nParameters are missing.\n");
	  fprintf(stdout, "Call with <ParameterFile>\n\n");
	}
      MPI_Finalize();
      exit(0);
    }

  read_parameterfile(argv[1]);

  checkchoose();  

  set_units();

  initialize_transferfunction(); 

  initialize_powerspectrum(); 

  initialize_ffts(); 

  read_glass(GlassFile);

  displacement_fields();

  if(NumPart)
    free(P);

  free_ffts();


  if(ThisTask == 0)
    {
      printf("\nDone!\n\n");
      printf("\n");
    }

  MPI_Barrier(MPI_COMM_WORLD);

  MPI_Finalize();		/* clean up & finalize MPI */
  exit(0);
}





void displacement_fields(void)
{
  gsl_rng *random_generator;
  int i, j, k, ii, jj;
  double fac, vel_prefac, vel_prefac2;
  double phig, Beta;
  double kvec[3], kmag, kmag2;
  double phase, ampl, hubble_a;
  double maxdisp;
  unsigned int *seedtable;
 
  unsigned int bytes, 
  nmesh3;
  int coord;

  fftw_complex *(cpot);  /* For computing nongaussian fnl ic */
  fftw_real *(pot);
  
  fftw_complex *(cppot); /* For two field case */
  fftw_real *(ppot);


#ifdef CORRECT_CIC
  double fx, fy, fz, ff, smth;
#endif



  if(ThisTask == 0)
    {
      printf("\nstart computing displacement fields...\n");
      fflush(stdout);
    }

  hubble_a =
    Hubble * sqrt(Omega / pow(InitTime, 3) + (1 - Omega - OmegaLambda) / pow(InitTime, 2) + OmegaLambda);

  vel_prefac = InitTime * hubble_a * F_Omega(InitTime);
  vel_prefac2 = InitTime * hubble_a * F2_Omega(InitTime);

  vel_prefac /= sqrt(InitTime);	/* converts to Gadget velocity */
  vel_prefac2 /= sqrt(InitTime);

  if(ThisTask == 0)
    printf("vel_prefac= %g, vel_prefac2= %g,  hubble_a=%g fom=%g \n", vel_prefac, vel_prefac2, 
                                                                      hubble_a, F_Omega(InitTime));

  fac = pow(2 * PI / Box, 1.5);

  maxdisp = 0;

  random_generator = gsl_rng_alloc(gsl_rng_ranlxd1);

  gsl_rng_set(random_generator, Seed);

  if(!(seedtable = malloc(Nmesh * Nmesh * sizeof(unsigned int))))
    FatalError(4);

  for(i = 0; i < Nmesh / 2; i++)
    {
      for(j = 0; j < i; j++)
	seedtable[i * Nmesh + j] = 0x7fffffff * gsl_rng_uniform(random_generator);

      for(j = 0; j < i + 1; j++)
	seedtable[j * Nmesh + i] = 0x7fffffff * gsl_rng_uniform(random_generator);

      for(j = 0; j < i; j++)
	seedtable[(Nmesh - 1 - i) * Nmesh + j] = 0x7fffffff * gsl_rng_uniform(random_generator);

      for(j = 0; j < i + 1; j++)
	seedtable[(Nmesh - 1 - j) * Nmesh + i] = 0x7fffffff * gsl_rng_uniform(random_generator);

      for(j = 0; j < i; j++)
	seedtable[i * Nmesh + (Nmesh - 1 - j)] = 0x7fffffff * gsl_rng_uniform(random_generator);

      for(j = 0; j < i + 1; j++)
	seedtable[j * Nmesh + (Nmesh - 1 - i)] = 0x7fffffff * gsl_rng_uniform(random_generator);

      for(j = 0; j < i; j++)
	seedtable[(Nmesh - 1 - i) * Nmesh + (Nmesh - 1 - j)] = 0x7fffffff * gsl_rng_uniform(random_generator);

      for(j = 0; j < i + 1; j++)
	seedtable[(Nmesh - 1 - j) * Nmesh + (Nmesh - 1 - i)] = 0x7fffffff * gsl_rng_uniform(random_generator);
    }



  /* non gaussian initial potential fnl type  */

      if(ThisTask == 0)
        {
          printf("\nstarting nongaussian with Fnl: %lf, qsig: %lf \n", Fnl, qsig);
          fflush(stdout);
        }


      bytes=0; /*initialize*/
      cpot = (fftw_complex *) malloc(bytes += sizeof(fftw_real) * TotalSizePlusAdditional);
      pot = (fftw_real *) cpot;

      ASSERT_ALLOC(cpot);

      /* first, clean the cpot array */
      for(i = 0; i < Local_nx; i++)
        for(j = 0; j < Nmesh; j++)
          for(k = 0; k <= Nmesh / 2; k++)
              {
                cpot[(i * Nmesh + j) * (Nmesh / 2 + 1) + k].re = 0;
                cpot[(i * Nmesh + j) * (Nmesh / 2 + 1) + k].im = 0;
              }


                                                                /* Ho in units of h/Mpc and c=1, i.e., internal units so far  */
      Beta = 1.5 * Omega / FnlTime / (2998. * 2998. );          /* Beta = 3/2 H(z)^2 a^2 Om(a) = 3/2 Ho^2 Om0 / a */ 


      for(i = 0; i < Nmesh; i++)
        {
          ii = Nmesh - i;
          if(ii == Nmesh)
            ii = 0;
          if((i >= Local_x_start && i < (Local_x_start + Local_nx)) ||
             (ii >= Local_x_start && ii < (Local_x_start + Local_nx)))
            {
              for(j = 0; j < Nmesh; j++)
                {
                  gsl_rng_set(random_generator, seedtable[i * Nmesh + j]);

                  for(k = 0; k < Nmesh / 2; k++)
                    {
                      phase = gsl_rng_uniform(random_generator) * 2* PI;
                      do
                        ampl = gsl_rng_uniform(random_generator);
                      while(ampl == 0);

                      if(i == Nmesh / 2 || j == Nmesh / 2 || k == Nmesh / 2)
                        continue;
                      if(i == 0 && j == 0 && k == 0)
                        continue;

                      if(i < Nmesh / 2)
                        kvec[0] = i * 2 * PI / Box;
                      else
                        kvec[0] = -(Nmesh - i) * 2 * PI / Box;

                      if(j < Nmesh / 2)
                        kvec[1] = j * 2 * PI / Box;
                      else
                        kvec[1] = -(Nmesh - j) * 2 * PI / Box;

                      if(k < Nmesh / 2)
                        kvec[2] = k * 2 * PI / Box;
                      else
                        kvec[2] = -(Nmesh - k) * 2 * PI / Box;

                      kmag2 = kvec[0] * kvec[0] + kvec[1] * kvec[1] + kvec[2] * kvec[2];
                      kmag = sqrt(kmag2);

                      if(SphereMode == 1)
                        {
                          if(kmag * Box / (2 * PI) > Nsample / 2)       /* select a sphere in k-space */
                            continue;
                        }
                      else
                        {
                          if(fabs(kvec[0]) * Box / (2 * PI) > Nsample / 2)
                            continue;
                          if(fabs(kvec[1]) * Box / (2 * PI) > Nsample / 2)
                            continue;
                          if(fabs(kvec[2]) * Box / (2 * PI) > Nsample / 2)
                            continue;
                        }

                      phig = -log(ampl) * Anorm * exp( 
                      PrimordialIndex * log(kmag) );   /* initial normalized power */
                      phig =  phig; /* for two field case, qsig needs to be defined */           
                      phig = sqrt(phig) * fac * Beta / DstartFnl / kmag2;    /* amplitude of the initial gaussian potential */
                      

                      if(k > 0)
                        {
                          if(i >= Local_x_start && i < (Local_x_start + Local_nx))
                               {

                                coord = ((i - Local_x_start) * Nmesh + j) * (Nmesh / 2 + 1) + k;

                                cpot[coord].re = phig * sin(phase);
                                cpot[coord].im = - phig * cos(phase);

                               }
                        }
                      else      /* k=0 plane needs special treatment */
                        {
                          if(i == 0)
                            {
                              if(j >= Nmesh / 2)
                                continue;
                              else
                                {
                                  if(i >= Local_x_start && i < (Local_x_start + Local_nx))
                                    {
                                      jj = Nmesh - j;   /* note: j!=0 surely holds at this point */

                                          coord = ((i - Local_x_start) * Nmesh + j) * (Nmesh / 2 + 1) + k;

                                          cpot[coord].re =  phig * sin(phase);
                                          cpot[coord].im = - phig * cos(phase);


                                          coord = ((i - Local_x_start) * Nmesh + jj) * (Nmesh / 2 + 1) + k; 
                                          cpot[coord].re = phig * sin(phase);
                                          cpot[coord].im = phig * cos(phase);

                                    }
                                }
                            }
                          else  /* here comes i!=0 : conjugate can be on other processor! */
                            {
                              if(i >= Nmesh / 2)
                                continue;
                              else
                                {
                                  ii = Nmesh - i;
                                  if(ii == Nmesh)
                                    ii = 0;
                                  jj = Nmesh - j;
                                  if(jj == Nmesh)
                                    jj = 0;

                                  if(i >= Local_x_start && i < (Local_x_start + Local_nx))
                                      {

                                        coord = ((i - Local_x_start) * Nmesh + j) * (Nmesh / 2 + 1) + k;
          
                                        cpot[coord].re = phig * sin(phase);
                                        cpot[coord].im = - phig * cos(phase);

                                      }
                                  if(ii >= Local_x_start && ii < (Local_x_start + Local_nx))
                                      {
                                        coord = ((ii - Local_x_start) * Nmesh + jj) * (Nmesh / 2 + 1) + k;

                                        cpot[coord].re = phig * sin(phase);
                                        cpot[coord].im = phig * cos(phase);

                                      }
                                }
                            }
                        }
                    }
                }
            }
        }

bytes=0;
cppot = (fftw_complex *) malloc(bytes += sizeof(fftw_real) * TotalSizePlusAdditional);
ppot = (fftw_real *) cppot;
ASSERT_ALLOC(cppot);

for (i=0; i< Local_nx; i++)
  for (j=0; j< Nmesh; j++) 
    for (k=0; k<=Nmesh/2; k++) {
	cppot[(i*Nmesh+j)*(Nmesh/2+1)+k].re = 0;
	cppot[(i*Nmesh+j)*(Nmesh/2+1)+k].im = 0;
}



     /*** For non-local models it is important to keep all factors of SQRT(-1) as done below ***/
     /*** Notice also that there is a minus to convert from Bardeen to gravitational potential ***/


      /******* LOCAL PRIMORDIAL POTENTIAL ************/

      if(ThisTask == 0) printf("Fourier transforming initial potential to configuration...");
      rfftwnd_mpi(Inverse_plan, 1, pot, Workspace, FFTW_NORMAL_ORDER);
      rfftwnd_mpi(Inverse_plan, 1, ppot, Workspace, FFTW_NORMAL_ORDER); 
      if(ThisTask == 0) printf("Done.\n");
      fflush(stdout);

      /* square the potential in configuration space and add the other gaussian 
field too! */
      double vart = 0; // analog of varphi in fnl_ic_mpi_2fld_redef.c
      double vartt=0;
      for(i = 0; i < Local_nx; i++)
        for(j = 0; j < Nmesh; j++)
          for(k = 0; k < Nmesh; k++)
            {
             coord = (i * Nmesh + j) * (2 * (Nmesh / 2 + 1)) + k; 
	 //    printf("psi: %g pot[coord], phi: %g ppot[coord]\n", pot[coord], ppot[coord]);
	     vart+=pot[coord]*pot[coord];
            }
      MPI_Allreduce(&vart, &vartt, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
      
      for(i = 0; i < Local_nx; i++)
        for(j = 0; j < Nmesh; j++)
          for(k = 0; k < Nmesh; k++)
            {
             coord = (i * Nmesh + j) * (2 * (Nmesh / 2 + 1)) + k; 
	 //    printf("psi: %g pot[coord], phi: %g ppot[coord]\n", pot[coord], ppot[coord]); 
	           ppot[coord]=(pot[coord]*pot[coord]-vartt);
            }
            
      MPI_Barrier(MPI_COMM_WORLD);

      if(ThisTask == 0) printf("Fourier transforming squared potential ...");
      //printf("vart: %f\n", vart);
      if (ThisTask==0) printf("varrr: %f\n", vartt);
      rfftwnd_mpi(Forward_plan, 1, pot, Workspace, FFTW_NORMAL_ORDER);
      rfftwnd_mpi(Forward_plan, 1, ppot, Workspace, FFTW_NORMAL_ORDER);
      if(ThisTask == 0) printf("Done.\n");
      fflush(stdout);
   
       /* remove the N^3 I got by forwardfurier and put zero to zero mode */
       /* also multiply by the smoothing function and the transfer function */
      nmesh3 = ((unsigned int) Nmesh) * ((unsigned int) Nmesh ) * ((unsigned int) Nmesh);    
      
      for(i = 0; i < Local_nx; i++)
        for(j = 0; j < Nmesh; j++)
          for(k = 0; k <= Nmesh / 2 ; k++)
            {
              coord = (i * Nmesh + j) * (Nmesh / 2 + 1) + k;
              cpot[coord].re /= (double) nmesh3; 
              cpot[coord].im /= (double) nmesh3; 
              cppot[coord].im /= (double) nmesh3;
              cppot[coord].re /= (double) nmesh3;
            }
        
       if(ThisTask == 0) {
              cpot[0].re=0.;
              cpot[0].im=0.; 
              cppot[0].re=0.;
              cppot[0].im=0.;
           }

      // generate log distribution of R 
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 4.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 6.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 8.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 10.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 12.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 14.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 5.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 7.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 9.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 11.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 13.);
      statistics(cpot, cppot, 1.0/(1+Redshift), Box, Nmesh, 15.);

MPI_Barrier(MPI_COMM_WORLD);
if (ThisTask==0) printf("finished..\n");
fflush(stdout);
// end here after getting the desired statistics 


    /****** FINISHED NON LOCAL POTENTIAL OR LOCAL FNL, STILL IN NONGAUSSIAN SECTION ****/   
    /*****  Now 2LPT ****/

}

float Wtophat(float x)
{
  if (fabs(x) < 0.03) return 1 - 0.1*x*x;
  return 3*(sin(x)-x*cos(x))/(x*x*x);
}


void statistics(fftw_complex *c1, fftw_complex *c2, float a_init, float xsize, long n1, float R)
{
  fftw_complex *Fs1 = NULL;
  fftw_real *(cFs1) = NULL;
    
  fftw_complex *Fs2 = NULL;
  fftw_real *(cFs2) = NULL;
    
  float k[3], ksq, dk=2*M_PI/xsize, sigma8sq = Sigma8*Sigma8, W;
  int ll, l, m, n, cord, bytes;
  double td1=0., td2=0.;
  double d1=0., d2=0., d2f=0., d3h=0., d3f=0., d4h=0., d4f=0., d5h=0., d5f=0., d3g=0., d4g=0., d5g=0.;
  double dd1=0., dd2=0., dd2f=0., dd3h=0., dd3f=0., dd4h=0., dd4f=0., dd5h=0., dd5f=0., dd3g=0., dd4g=0., dd5g=0.;
  double sum=0., variance=0.;
  //rfftwnd_mpi_plan plan;
  int local_nx, local_x_start, local_ny_after_transpose,
    local_y_start_after_transpose, total_local_size;
 rfftwnd_mpi_local_sizes(Inverse_plan, &local_nx, &local_x_start,
			  &local_ny_after_transpose, &local_y_start_after_transpose, &total_local_size);
  //plan = rfftw3d_mpi_create_plan(MPI_COMM_WORLD, n1,n1,n1, FFTW_COMPLEX_TO_REAL, FFTW_ESTIMATE);  // this is Inverse_plan
  bytes=0;
  Fs1 = (fftw_complex *) malloc(bytes += sizeof(fftw_real) * 
TotalSizePlusAdditional);
  cFs1 = (fftw_real *) Fs1;

  bytes=0;
  Fs2 = (fftw_complex *) malloc(bytes += sizeof(fftw_real) * 
TotalSizePlusAdditional);
  cFs2 = (fftw_real *) Fs2;
  
 // printf("entering the loop with %d %d\n", local_ny_after_transpose, (int) n1);
double Beta = 1.5 * Omega / FnlTime / (2998. * 2998. );          /* Beta = 3/2 H(z)^2 a^2 Om(a) = 3/2 Ho^2 Om0 / a */ 

  for (ll=0; ll<local_ny_after_transpose; ll++) {
    l = ll + local_y_start_after_transpose;
    k[1] = dk * (l < (n1/2 + 1) ? l : l-n1);
    for (m=0; m<n1; m++) {
      k[0] = dk * (m < (n1/2 + 1) ? m : m-n1);
      for (n=0; n < (n1/2 + 1); n++) {
	k[2] = dk*n;
	ksq = dot(k,k);
	W = Wtophat(R*sqrt(ksq));
	cord = (ll * n1 + m) * ((n1 / 2 + 1)) + n;  
	Fs1[cord].re = c1[cord].re * W * DstartFnl*ksq*TransferFunc(sqrt(ksq))/Beta;
	Fs1[cord].im = c1[cord].im * W * DstartFnl*ksq*TransferFunc(sqrt(ksq))/Beta;
	Fs2[cord].re = c2[cord].re* W * DstartFnl*ksq*TransferFunc(sqrt(ksq))/Beta;
	Fs2[cord].im = c2[cord].im * W * DstartFnl*ksq*TransferFunc(sqrt(ksq))/Beta; 
//	printf("%f %f %d %g\n", ksq, W, cord, c1[cord].re);
      }
    }
  }

  rfftwnd_mpi(Inverse_plan, 1, cFs1, Workspace, FFTW_NORMAL_ORDER);
  rfftwnd_mpi(Inverse_plan, 1, cFs2, Workspace, FFTW_NORMAL_ORDER);
  //rfftwnd_mpi_destroy_plan(plan);
 
 // now generate the statiscs
  for (l=0; l<local_nx*n1; l++) { 
    for (m=0; m<n1; m++) {
      dd1 = cFs1[l*(n1+2)+m];
      dd2 = cFs2[l*(n1+2)+m];
      td1 += dd1;
      td2 += dd2;
      sum += sqr(dd1);
      dd2f += sqr(dd2);
      dd3g += pow(dd1,3.0);
      dd3h += 3*sqr(dd1)*dd2;
      dd3f += pow(dd2, 3.0);
      dd4g += pow(dd1,4.0);
      dd4h += 6*sqr(dd2)*sqr(dd1);
      dd4f += pow(dd2,4.0);
      dd5g += pow(dd1,5.0);
      dd5h += 10*pow(dd2,3.0)*sqr(dd1);
      dd5f += pow(dd2,5.0);
//printf("%lf %lf %g\n", dd1, sum, dd3h);
    }
  }
  
  //printf("sum: %lf \n", sum);
  MPI_Allreduce(&sum, &variance, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  MPI_Allreduce(&td1, &d1, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  MPI_Allreduce(&td2, &d2, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  MPI_Allreduce(&dd2f, &d2f, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  MPI_Allreduce(&dd3g, &d3g, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  MPI_Allreduce(&dd3h, &d3h, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);  
  MPI_Allreduce(&dd3f, &d3f, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  MPI_Allreduce(&dd4g, &d4g, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);  
  MPI_Allreduce(&dd4h, &d4h, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);  
  MPI_Allreduce(&dd4f, &d4f, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  MPI_Allreduce(&dd5g, &d5g, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);  
  MPI_Allreduce(&dd5h, &d5h, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  MPI_Allreduce(&dd5f, &d5f, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  double n13=n1*n1*n1;
  variance /= n13; d1/=n13; d2/=n13; d2f/=n13; d3h/=n13; d3f/=n13; d4h/=n13; d4f/=n13; d5h/=n13; d5f/=n13, d3g/=n13, d4g/=n13, d5g/=n13;
  printf("var: %g, exp: %g\n", variance, sigma8sq);

  FILE* fpdata;
  char fname[100];
  sprintf(fname, "/gpfs/home/sza5154/group/moments/data%f_%d.txt", R, Seed);
  fpdata = fopen(fname, "w");

  fprintf(fpdata, "11 %g\n12 %g\n21 %g\n22 %g\n31 %g\n32 %g\n41 %g\n42 %g\n51 %g\n52 %g\n30 %g\n40 %g\n50 %g", d1, d2, variance, d2f, d3h, d3f, d4h, d4f, d5h, d5f, d3g, d4g, d5g);
  fclose(fpdata);

 free(Fs1);free(Fs2);

  return;
}

double periodic_wrap(double x)
{
  while(x >= Box)
    x -= Box;

  while(x < 0)
    x += Box;

  return x;
}


void set_units(void)		/* ... set some units */
{
  UnitTime_in_s = UnitLength_in_cm / UnitVelocity_in_cm_per_s;

  G = GRAVITY / pow(UnitLength_in_cm, 3) * UnitMass_in_g * pow(UnitTime_in_s, 2);
  Hubble = HUBBLE * UnitTime_in_s;
}



void initialize_ffts(void)
{
  int total_size, i, additional;
  int local_ny_after_transpose, local_y_start_after_transpose;
  int *slab_to_task_local;
  size_t bytes;


  Inverse_plan = rfftw3d_mpi_create_plan(MPI_COMM_WORLD,
					 Nmesh, Nmesh, Nmesh, FFTW_COMPLEX_TO_REAL, FFTW_ESTIMATE);

  Forward_plan = rfftw3d_mpi_create_plan(MPI_COMM_WORLD,
					 Nmesh, Nmesh, Nmesh, FFTW_REAL_TO_COMPLEX, FFTW_ESTIMATE);

  rfftwnd_mpi_local_sizes(Forward_plan, &Local_nx, &Local_x_start,
			  &local_ny_after_transpose, &local_y_start_after_transpose, &total_size);

  Local_nx_table = malloc(sizeof(int) * NTask);
  MPI_Allgather(&Local_nx, 1, MPI_INT, Local_nx_table, 1, MPI_INT, MPI_COMM_WORLD);

  if(ThisTask == 0)
    {
      for(i = 0; i < NTask; i++)
	printf("Task=%d Local_nx=%d\n", i, Local_nx_table[i]);
      fflush(stdout);
    }


  Slab_to_task = malloc(sizeof(int) * Nmesh);
  slab_to_task_local = malloc(sizeof(int) * Nmesh);

  for(i = 0; i < Nmesh; i++)
    slab_to_task_local[i] = 0;

  for(i = 0; i < Local_nx; i++)
    slab_to_task_local[Local_x_start + i] = ThisTask;

  MPI_Allreduce(slab_to_task_local, Slab_to_task, Nmesh, MPI_INT, MPI_SUM, MPI_COMM_WORLD);

  free(slab_to_task_local);



  additional = (Nmesh) * (2 * (Nmesh / 2 + 1));	/* additional plane on the right side */

  TotalSizePlusAdditional = total_size + additional;

  Workspace = (fftw_real *) malloc(bytes = sizeof(fftw_real) * total_size);

  ASSERT_ALLOC(Workspace)

}



void free_ffts(void)
{
  free(Workspace);
  free(Slab_to_task);
  rfftwnd_mpi_destroy_plan(Inverse_plan);
  rfftwnd_mpi_destroy_plan(Forward_plan);
}


int FatalError(int errnum)
{
  printf("FatalError called with number=%d\n", errnum);
  fflush(stdout);
  MPI_Abort(MPI_COMM_WORLD, errnum);
  exit(0);
}


